﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;

namespace WardenMonitor
{

    public class CheckTypeGuess
    {
        public CheckType Type { get; private set; }
        public int Confidence { get; set; }

        public CheckTypeGuess(CheckType type)
        {
            Type = type;
            Confidence = 1;
        }
    }

    public class HeuristicEngine
    {
        private const int MinimumIdentityConfidence = 10;
        private const int MinimumConfidence = 35;
        private const int UpperLengthLimit = 64;

        private int _processedIdentityPackets;
        private int _consecutiveKnown;
        private bool _isConfident;
        private bool _isIdentified;

        public List<byte> KnownCheckTypes = new List<byte>();
        public Dictionary<byte, List<CheckTypeGuess>> TypeGuesses = new Dictionary<byte, List<CheckTypeGuess>>();

        private static byte ToCheckType(byte raw, CheatCheckPacket packet)
        {
            return (byte) (raw ^ packet.Xor);
        }
        
        public HeuristicEngine(PacketHandler handler)
        {
            handler.OnIncomingPacket += OnPacket;
        }

        private void OnPacket(object sender, PacketEvent e)
        {
            if (e.Packet.Opcode == 0x02)
                HandleCheatCheck(new CheatCheckPacket(e.Packet.Data));
        }

        private void HandleCheatCheck(CheatCheckPacket packet)
        {
            // step 0: gather all check types by collecting the first byte after strings which is guaranteed to be a check type
            if (!_isConfident)
            {
                var checkType = ToCheckType(packet.Data[packet.FirstCheckIndex], packet);
                if (!KnownCheckTypes.Contains(checkType))
                {
                    Log.WriteLine(DateTime.Now + " Adding CheckType: " + checkType.ToString("X2"));
                    KnownCheckTypes.Add(checkType);
                    _consecutiveKnown = 0;
                }
                else
                    _consecutiveKnown++;
                
                // if we've a string, we can do some extra heuristics
                if (packet.Strings.Count > 1)
                {
                    // minus two, xor byte and the one byte we look ahead
                    int to = packet.Data.Length - 2;
                    for (int i = 27; i < to; i++)
                    {
                        // a valid string index
                        if (packet.Data[i] == 0x01)
                        {
                            // if the next byte is an already known type, we can go back and add the type of this check as known
                            if (KnownCheckTypes.Contains(ToCheckType(packet.Data[i + 1], packet)))
                            {
                                var drvCheckType = ToCheckType(packet.Data[i - 25], packet);
                                if (KnownCheckTypes.Contains(drvCheckType))
                                    break;
                                KnownCheckTypes.Add(drvCheckType);
                                Log.WriteLine(
                                    DateTime.Now +
                                    " String Heuristic guesses {0:X2} as DriverCheck type. Adding to known check types.",
                                    drvCheckType);
                            }
                        }
                    }
                }

                if (_consecutiveKnown >= MinimumConfidence)
                {
                    Log.WriteLine(DateTime.Now + " Finished collecting check types. Known types:");
                    foreach (byte ct in KnownCheckTypes)
                        Log.WriteLine("\t{0:X2}", ct);

                    _isConfident = true;
                }
                else
                    return;
            }

            // step 1: guess actual types
            if (!_isIdentified)
            {
                int index = packet.FirstCheckIndex;
                while (true)
                {
                    index = GuessType(packet, index);
                    if (index == 0)
                        break;
                }
                _processedIdentityPackets++;

                if (_processedIdentityPackets >= MinimumIdentityConfidence)
                {
                    Log.WriteLine(DateTime.Now + " Finished guessing identity. Results:");
                    foreach (byte ct in TypeGuesses.Keys)
                    {
                        Log.WriteLine("\tGuesses for {0:X2}:", ct);
                        var guesses = TypeGuesses[ct].OrderByDescending(g => g.Confidence);
                        foreach (var guess in guesses)
                            Log.WriteLine("\t\t" + guess.Type + " Confidence: " + guess.Confidence);
                    }
                    _isIdentified = true;
                }
                else
                    return;
            }

            try
            {
                // step 2: actually parse the packet, and don't include the xor byte
                var reader = new BinaryReader(new MemoryStream(packet.Data, 0, packet.Data.Length - 1));
                // skip strings
                reader.ReadBytes(packet.FirstCheckIndex);
                while (true)
                {
                    var bCheckType = ToCheckType(reader.ReadByte(), packet);
                    var checkType = GetGuessedCheckType(bCheckType);
                    if (checkType == CheckType.Invalid)
                    {
                        Log.WriteLine(DateTime.Now + " Unknown CheckType " + bCheckType.ToString("X2"));
                        break;
                    }

                    switch (checkType)
                    {
                        case CheckType.MemoryCheck:
                            {
                                byte stringIndex = reader.ReadByte();
                                uint address = reader.ReadUInt32();
                                byte length = reader.ReadByte();
                                Log.WriteLine(DateTime.Now + " MemoryCheck: " + length + " bytes from 0x" +
                                              address.ToString("X8") + " Module: " +
                                              (stringIndex == 0 ? "default" : packet.Strings[stringIndex]));
                                break;
                            }
                        case CheckType.PageCheckA:
                            {
                                uint seed = reader.ReadUInt32();
                                var hash = reader.ReadBytes(20);
                                uint address = reader.ReadUInt32();
                                byte length = reader.ReadByte();
                                Log.WriteLine(DateTime.Now + " PageCheck: " + length + " bytes from 0x" +
                                              address.ToString("X8") + " Seed: " + seed.ToString("X8") + " Hash: " +
                                              BitConverter.ToString(hash));
                                break;
                            }
                        case CheckType.DriverCheck:
                            {
                                uint seed = reader.ReadUInt32();
                                var hash = reader.ReadBytes(20);
                                byte stringIndex = reader.ReadByte();
                                Log.WriteLine(DateTime.Now + " DriverCheck: " + packet.Strings[stringIndex] + " Seed: " +
                                              seed.ToString("X8") + " Hash: " + BitConverter.ToString(hash));
                                break;
                            }
                        default:
                            {
                                Log.WriteLine(DateTime.Now + " Can't handle " + checkType);
                                return;
                            }
                    }
                }
            }
            catch (Exception e)
            {
                Log.WriteLine(e.ToString());
            }

            Log.WriteLine();
        }

        private CheckType GetGuessedCheckType(byte ct)
        {
            if (!TypeGuesses.ContainsKey(ct))
                return CheckType.Invalid;

            return TypeGuesses[ct].OrderByDescending(g => g.Confidence).First().Type;
        }

        private int GuessType(CheatCheckPacket packet, int index)
        {
            if (index >= packet.Data.Length)
                return 0;

            var checkType = ToCheckType(packet.Data[index], packet);

            if (CheckMemoryCheck(packet, index))
            {
                AddGuess(checkType, CheckType.MemoryCheck);
                return index + 7;
            }

            if (CheckPageCheck(packet, index))
            {
                AddGuess(checkType, CheckType.PageCheckA);
                return index + 30;
            }

            if (packet.Strings.Count > 1)
            {
                if (CheckDriverCheck(packet, index))
                {
                    AddGuess(checkType, CheckType.DriverCheck);
                    return index + 26;
                }

                if (CheckProcCheck(packet, index))
                {
                    AddGuess(checkType, CheckType.ProcedureCheck);
                    return index + 32;
                }

                if (CheckMpqLuaCheck(packet, index))
                {
                    AddGuess(checkType, CheckType.MpqCheck);
                    AddGuess(checkType, CheckType.LuaStringCheck);
                }
            }

            return 0;
        }

        private void AddGuess(byte check, CheckType guess)
        {
            if (!TypeGuesses.ContainsKey(check))
                TypeGuesses.Add(check, new List<CheckTypeGuess>());

            var list = TypeGuesses[check];
            if (list.Any(g => g.Type == guess))
                list.Where(g => g.Type == guess).First().Confidence++;
            else
                list.Add(new CheckTypeGuess(guess));
        }

        // TODO: not seen in a while. needs to evaluate string to differentiate between MPQ and LUA check
        private bool CheckMpqLuaCheck(CheatCheckPacket packet, int index)
        {
            if (!CheckEnd(packet, index, 1))
                return false;
            
            // string index to indicate LUA string or MPQ file
            byte stringIndex = packet.Data[index + 1];
            if (stringIndex == 0 || stringIndex > packet.Strings.Count)
                return false;

            return true;
        }

        // TODO: not seen for some time now. can string indexes be 0?
        private bool CheckProcCheck(CheatCheckPacket packet, int index)
        {
            if (!CheckEnd(packet, index, 31))
                return false;

            // two string indexes
            byte stringIndex = packet.Data[index + 25];
            if (stringIndex == 0 || stringIndex > packet.Strings.Count)
                return false;
            stringIndex = packet.Data[index + 26];
            if (stringIndex == 0 || stringIndex > packet.Strings.Count)
                return false;

            // and a length at the end
            if (packet.Data[index + 31] > UpperLengthLimit)
                return false;

            return true;
        }

        private bool CheckDriverCheck(CheatCheckPacket packet, int index)
        {
            if (!CheckEnd(packet, index, 25))
                return false;

            // string index for driver name, no default value (= can not be null)
            byte stringIndex = packet.Data[index + 25];
            if (stringIndex == 0 || stringIndex > packet.Strings.Count)
                return false;

            return true;
        }

        private bool CheckPageCheck(CheatCheckPacket packet, int index)
        {
            if (!CheckEnd(packet, index, 29))
                return false;

            // length at end
            if (packet.Data[index + 29] > UpperLengthLimit)
                return false;

            return true;
        }

        private bool CheckMemoryCheck(CheatCheckPacket packet, int index)
        {
            if (!CheckEnd(packet, index, 6))
                return false;

            // string index can be zero or a valid index into the string list
            byte stringIndex = packet.Data[index + 1];
            if (stringIndex != 0)
            {
                if (stringIndex > packet.Strings.Count)
                    return false;
            }

            // next up a valid memory address
            if (!ValidateAddress(packet, index + 2))
                return false;

            // and a reasonable length
            if (packet.Data[index + 6] > UpperLengthLimit)
                return false;

            return true;
        }

        // very basic memory heuristic. this can probably be refined by actually checking memory.
        private static bool ValidateAddress(CheatCheckPacket packet, int index)
        {
            return packet.Data[index + 3] == 0;
        }

        private bool CheckEnd(CheatCheckPacket packet, int startIndex, int size)
        {
            int offset = startIndex + size + 1;
            // note: EOF, but size matches perfectly - so report true
            if (offset == packet.Data.Length)
                return true;
            // note: EOF, and size exceeds length - impossible
            if (offset > packet.Data.Length)
                return false;
            return KnownCheckTypes.Contains(ToCheckType(packet.Data[offset], packet));
        }
    }

}